import requests
from django.shortcuts import render
from langchain.chains.llm import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain_community.document_loaders import TextLoader
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import PromptTemplate
from rest_framework.response import Response
from rest_framework.views import APIView
from langchain_text_splitters import CharacterTextSplitter
from tools.faiss import faissdb


# Create your views here.


class GetData(APIView):
    def get(self, request):
        doc = TextLoader("AI/doc/doctor.txt", encoding='utf-8').load()
        spliter = CharacterTextSplitter("\n", chunk_size=20, chunk_overlap=5)
        chunks = spliter.split_documents(doc)
        faissdb.add(chunks, 'ask')
        return Response({"msg": "ok"})


class GetFaissMes(APIView):
    # 初始化类级别的 memory
    memory = ConversationBufferMemory()

    def get(self, request):
        # 获取用户的问题
        ask = request.GET.get("ask")
        # 根据问题查询相似度最高的三个问题
        try:
            askres = faissdb.search(ask, 'ask', 3)
            # 提取 askres 中的文本内容，确保输入给 dashscope 的是文本而不是 Document 对象
            askres_text = [item.page_content if hasattr(item, 'page_content') else str(item) for item in askres]
            # 查询三个问题的答案
            res = faissdb.search(askres_text, 'movies', 3)
        except Exception as e:
            return Response({"error": str(e)}, status=500)

        llm = Tongyi()
        # 定义 Prompt 模板
        template = """{input}
                 """
        prompt = PromptTemplate(input_variables=["input"], template=template)
        # 初始化 LLMChain
        self.memory.chat_memory.add_user_message(ask)
        # 将用户输入和 AI 响应添加到对话历史记录中
        inputmes = "问题是" + ask + ",答案为" + str(res)
        chain = LLMChain(llm=llm, prompt=prompt, memory=self.memory)
        aires = chain.run(input=inputmes)
        self.memory.chat_memory.add_ai_message(aires)

        # 解析 memory，取最后三轮
        prompt = "内容为{mes},从以上内容中提取所有的问题和答案，帮我提取关键信息总结，返回一条问答和一条答案"
        promptTemplate = PromptTemplate.from_template(prompt)
        # 生成 prompt
        prompt = promptTemplate.format(mes=str(self.memory))
        # 实例化通义大模型
        tongyi = Tongyi()
        ret = tongyi.invoke(prompt)
        print(ret)
        # 清空 memory 而不是重新初始化
        self.memory.clear()
        # 根据 ret 的实际类型修改以下逻辑
        # 假设 ret 是一个字符串，你可能需要使用其他逻辑而不是 for 循环
        # 以下是一个可能的修改示例，假设 ret 是一个字符串
        self.memory.chat_memory.add_user_message(ret)
        self.memory.chat_memory.add_ai_message(aires)

        return Response({"code": 200, 'aires': aires})
